# coding: utf8
"""
client:
python icmp_tun.py 10.0.0.1 10.21.210.237

server:
python icmp_tun.py 10.0.0.2

client -----> server
       <-----
       <====>
"""


import sys
import os
import optparse
import socket
import select
import errno
import struct
import random

from collections import OrderedDict
import pytun
from idlelib.AutoComplete import ID_CHARS

_MAGIC_ID = "XINGICMPTUNNEL"
_MAGIC_LEN = len(_MAGIC_ID)

def icmp_checksum(source_string):
    """
    I'm not too confident that this is right but testing seems
    to suggest that it gives the same answers as in_cksum in ping.c
    """

    # return random.randint(0, 0xFFFF)
    sum_ = 0
    count_to = (len(source_string) / 2) * 2
    for count in xrange(0, count_to, 2):
        this = ord(source_string[count + 1]) * 256 + ord(source_string[count])
        sum_ += this
        sum_ &= 0xffffffff  # Necessary?

    if count_to < len(source_string):
        sum_ += ord(source_string[len(source_string) - 1])
        sum_ &= 0xffffffff  # Necessary?

    sum_ = (sum_ >> 16) + (sum_ & 0xffff)
    sum_ += (sum_ >> 16)
    answer = ~sum_
    answer &= 0xffff

    # Swap bytes. Bugger me if I know why.
    answer = answer >> 8 | (answer << 8 & 0xff00)

    return answer


def encrypt(data):
    return data

def decrypt(data):
    return data

class ICMPPocket(object):
    def __init__(self, type_, id_, seq, data, checksum=0, code=0, addr=None):
        self.type = type_
        self.code = code
        self.checksum = checksum
        self.id = id_
        self.seq = seq
        self.data = data
        self.addr = addr

    def __str__(self):
        return "addr: %s, type:%d, id:%d, seq:%d, data:%s ..." % \
            (self.addr, self.type, self.id, self.seq, str(self.data[:64]))

    @classmethod
    def parse(cls, obj, recvlen=None):
        if recvlen is not None:
            # socket
            raw, addr = obj.recvfrom(recvlen)
            addr = addr[0]
        else:
            raw = obj
            addr = None
        type_, code, checksum, id_, seq = struct.unpack("!BBHHH", raw[20:28])
        if raw[28:28+_MAGIC_LEN] != _MAGIC_ID:
            print "MAGIC_ID error"
            return None
        data = decrypt(raw[28+_MAGIC_LEN:])
        return cls(type_, id_, seq, data, checksum, code, addr)

    def pack(self):
        self.data = _MAGIC_ID + self.data
        packfmt = "!BBHHH%ds" % (len(self.data))
        # checksum = random.randint(0, 0xFFFF)  # save cpu
        data = encrypt(self.data)
        args = [self.type, self.code, 0, self.id, self.seq, data]
        args[2] = icmp_checksum(struct.pack(packfmt, *args))
        return struct.pack(packfmt, *args)

    def sendto(self, sock, host):
        sock.sendto(self.pack(), (host, 0))
    
    
class Connection(object):
    def __init__(self, id_, seq, icmp_tun, tcp_sock, icmp_sock):
        self._id = id_
        sefl._icmp_tun = icmp_tun
        self._tcp_sock = tcp_sock
        self._icmp_sock = icmp_sock
        self._mtu = self._icmp_tun._tun.mtu
        self._seq = seq
        
    def _generate_seq(self):
        self._seq += 2
        self._seq &= 0xFFFF
        return self._seq
        
#     def write_tun(self, data):
#         return self._tun.write(data)
#     
#     def read_tun(self):
#         return self._tun.read(self._mtu)
    
    def write_tcp(self, data):
        return self._tcp_sock.send(data)
    
    def read_tcp(self, data):
        return self._tcp_sock.recv(self._mtu)
    
    def write_icmp(self, data):
        ICMPPocket(self._icmp_type, self._icmp_id, self.generate_seq(), to_sock
                                    ).sendto(self._icmp_sock, self._peer_addr)
    
    def read_icmp(self):
        pass
#         icmp_pocket = ICMPPocket.parse(self._icmp_sock, self._mtu)
#         if icmp_pocket:
#             if icmp_pocket.raddr == self._peer_addr and \
#                 icmp_pocket.seq&0x01 != self._seq&0x01:  # 不是自己发出的包
#                 if self._icmp_id is None:
#                     self._icmp_id = icmp_pocket.id
#                 else:
#                     to_tun = icmp_pocket.data
#                 print "recv", len(to_tun), "from", raddr

class BaseTunnel(object):
    def __init__(self, lt_addr, pt_addr=None, mtu=1200):
        self._tun = pytun.TunTapDevice(flags=pytun.IFF_TUN|pytun.IFF_NO_PI)
        self._tun.addr = lt_addr
        self._tun.netmask = '255.255.255.0'
        self._tun.mtu = mtu
        self._tun.up()
        
        self.lt_addr = lt_addr
        
        if pt_addr is not None:
            self.pt_addr = pt_addr
        else:
            ip_split = map(int, lt_addr.split("."))
            ip_split[3] = (ip_split[3]+1)%255
            self.pt_addr = ".".join(map(str, ip_split))
            
        self._seq = None    # icmp seq
            
    def generate_seq(self):
        self._seq += 2
        self._seq &= 0xFFFF
        return self._seq

class Tunnel(object):

    def __init__(self, taddr, peer_addr, dest_addr=None, dest_port=None, lport=None):
        self._tun = pytun.TunTapDevice(flags=pytun.IFF_TUN|pytun.IFF_NO_PI)
        self._tun.addr = taddr
        self._tun.netmask = '255.255.255.0'
        self._tun.mtu = 1200
        self._tun.up()
        self._peer_addr = peer_addr

        self._icmp_sock = socket.socket(socket.AF_INET, socket.SOCK_RAW,
						    socket.getprotobyname("icmp"))

        self._is_client = lport is not None
        if self._is_client:
            self._icmp_id = 0xFFFF & os.getpid()
            self._seq = 0
            self._icmp_type = 8

            self._lport = lport
            self._tcp_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 
                                socket.getprotobyname("tcp"))
            self._tcp_sock.bind("0.0.0.0", lport)

            print "tpc:%d <==> %s(%s) <== icmp ==> %s" % (
                    lport, taddr, self._tun.name, peer_addr)
        else:
            self._icmp_id = None
            self._seq = 1
            self._icmp_type = 0

            self._dest_addr = dest_addr
            self._dest_port = dest_port

            print "(%s) <== icmp ==> %s(%s) <==> %s:%d" % (
                    peer_addr, taddr, self._tun.name, dest_addr, dest_port)            
        
    def generate_seq(self):
        self._seq += 2
        self._seq &= 0xFFFF
        return self._seq

    def run(self):
        mtu = self._tun.mtu
        to_tun = ''
        to_sock = ''
        raddr = None
        r = [self._tun, self._icmp_sock, self._tcp_sock]; w = []; x = []
        while 1:
            try:
                r, w, x = select.select(r, w, x)
                if self._tun in r:
                    to_sock = self._tun.read(mtu)

                if self._icmp_sock in r:
                    old_addr = raddr
                    icmp_pocket = ICMPPocket.parse(self._icmp_sock, 65535)
                    if icmp_pocket:
                        if icmp_pocket.raddr == self._peer_addr and \
                            icmp_pocket.seq&0x01 != self._seq&0x01:  # 不是自己发出的包
                            if self._icmp_id is None:
                                self._icmp_id = icmp_pocket.id
                            else:
                                to_tun = icmp_pocket.data
                                print "recv", len(to_tun), "from", raddr

                if self._tun in w:
                    self._tun.write(to_tun)
                    to_tun = ''

                if self._icmp_sock in w:
                    print "send ", len(to_sock), "to:", self._peer_addr
                    ICMPPocket(self._icmp_type, self._icmp_id, self.generate_seq(), to_sock
                                    ).sendto(self._icmp_sock, self._peer_addr)
                    to_sock = ''

                r = []; w = []
                if to_tun:
                    w.append(self._tun)
                else:
                    r.append(self._icmp_sock)
                if to_sock:
                    w.append(self._icmp_sock)
                else:
                    r.append(self._tun)

            except (select.error, socket.error, pytun.Error), e:
                if e[0] == errno.EINTR:
                    continue
                print >> sys.stderr, str(e)
                break

def main():
    try:
        # def __init__(self, taddr, peer_addr, dest_addr=None, dest_port=None, lport=None):
        if len(sys.argv) < 3:
            print "# client: ./tunnel 10.0.0.1 2.2.2.2 1234\n" \
                  "# proxy: ./tunnel 10.0.0.2 1.1.1.1 127.0.0.1:80"
            return 1

        if sys.argv[3].find(":") >= 0: # proxy
            daddr, dport = sys.argv[3].split(":")
            server = Tunnel(sys.argv[1], sys.argv[2], dest_addr=daddr, dest_port=int(dport))
        else: # client
            server = Tunnel(sys.argv[1], sys.argv[2], lport=int(sys.argv[3]))

    except (pytun.Error, socket.error), e:
        print >> sys.stderr, str(e)
        return 1
    
    try:
        server.run()
    except Exception, e:
        server._tun.down()
        server._tun.close()
        import tracback
        tracback.print_exc()
        return 3
    return 0

if __name__ == '__main__':
    sys.exit(main())

